import numpy as np
import os
import output_function as of
import pandas as pd
import matplotlib.pyplot as plt

# base_path = './data/RB_Ra=20000_Lx=5.0/'  # 根据需要修改为您的保存路径
# out_iter =30000  # 假设我们要加载第300时间步的数据

# base_path = './data/RB_Ra=4000_Lx=5.0/'  # 根据需要修改为您的保存路径
# out_iter =40000  # 假设我们要加载第300时间步的数据

base_path = './data/RB_Ra=120000_Lx=5.0_conserve/'  # 根据需要修改为您的保存路径
out_iter =14000  # 假设我们要加载第300时间步的数据

# 加载数据
n, psi, omega, u, v, out_iter_loaded = of.load_simulation_data(base_path, out_iter)
n_init0_matrix, xx, yy,x,y = of.load_initial_data(load_path=f'{base_path}initial_conditions.npz')

n_density=n

# 加载模拟参数
config = of.load_simulation_config(f'{base_path}simulation_config.json')
# 从config字典中提取特定的参数并赋值给变量
dx = config['dx']
dy = config['dy']
Lx = config['Lx']
Ly = config['Ly']
Nx = config['Nx']
Ny = config['Ny']
n_up = config['n_up']
n_0 = config['n_0']
n_down = config['n_down']
Delta_n = config['Delta_n']
Ra=Ra_star = config['Ra_star']
prandtl=Pr = config['Pr']
dt = config['dt']
ntime = config['ntime']
ndiag = config['ndiag']

# 打印一些变量以确认它们已被正确加载和赋值
print(f"dx: {dx}")
print(f"Ra_star: {Ra_star}")
print(f"dt: {dt}")
# 确认是否成功加载了正确的时间步长
if out_iter_loaded == out_iter:
    print(f"Successfully loaded data for time step: {out_iter_loaded}")
    print(omega.shape)
else:
    print("Failed to load the correct time step or file does not exist.")

# 接下来，您可以使用加载的数据进行分析、绘图或其他处理
# 接下来，您可以使用加载的数据进行分析、绘图或其他处理
#of.diag_n_pcolor(n,n_init0_matrix,u,v,xx,yy,out_iter,Ra_star,dt)
#of.diag_n_origin_pcolor_filepath(n,u,v,xx,yy,out_iter,Ra_star,dt)
t_total=out_iter*dt #总时间

# 初始化n_ave数组，用于存储每个j值对应的n的平均值
n_ave = np.zeros(Ny-1)  # Ny为n数组在y轴方向的长度
n_init_ave=np.zeros(Ny-1)
convec_n=np.zeros(Ny-1)
diff_n=np.zeros(Ny-1)
total_flux=np.zeros(Ny-1)
radial_x=np.zeros(Ny-1)
# 对于每个j值，计算对应列的平均值
for j in range(Ny-1):
    n_ave[j] = np.mean(n[j, :])
    n_init_ave[j]=np.mean(n_init0_matrix[j, :])
    radial_x[j]=y[j]
    for i in range(Nx-1):
        convec_n[j]+=n[j,i]*v[j,i]*dx
        diff_n[j]+=(n[j,i]-n[j+1,i])/dy*dx
        total_flux[j]=convec_n[j]+diff_n[j]


# 绘制n_ave与y的关系图
# plt.figure(figsize=(10, 6))
# plt.plot(radial_x, n_ave, label='Convective Density Density', color='red')
# plt.plot(radial_x, n_init_ave, label='Diffusive Density profile', color='blue')

# plt.axvline(x=Ly/2, color='black', linestyle='--', linewidth=1, label='y midplane')
# plt.xlabel('y')
# plt.ylabel('Average T ')
# plt.title(f'Convective vs. Diffusive T profile,Ra={Ra_star}')
# plt.legend()
# plt.grid(True)
# plt.show()

# 绘制n_ave与y的关系图
plt.figure(figsize=(10, 4))
plt.plot(radial_x, convec_n/Ly, label=r'$<vT>$', color='red')
plt.plot(radial_x, diff_n/Ly, label=r'$-<\partial T/\partial y>$', color='blue')
plt.plot(radial_x, total_flux/Ly, label=r'$<vT>-<\partial T/\partial y>$',color='purple')
plt.axvline(x=0.5, color='black', linestyle='--', linewidth=1, label='y-midplane')
plt.xlabel('position y')
plt.ylabel('Transfer flux')
plt.title(f'Convective vs. Diffusive transport,Ra={Ra_star},t={t_total:.3f}')
plt.legend()
plt.grid(True)
plt.show()



t_total=out_iter*dt #总时间
flag_pcolormesh_stream=True
if flag_pcolormesh_stream:
    t_diag=out_iter*dt
    fig, ax = plt.subplots(figsize=(10, 3))
    n_difference=n_density
    # 绘制伪彩色图
    
    plt.streamplot(xx, yy, u, v, linewidth=1)
    
    # plt.pcolormesh(xx, yy, n_difference, cmap=plt.get_cmap('Spectral_r'), vmin=-np.abs(n_difference).max(), vmax=np.abs(n_difference).max(), shading='auto')
    # plt.colorbar(label='n',boundaries=np.linspace(-np.abs(n_difference).max(), np.abs(n_difference).max()))
    
    # 使用自动缩放的颜色条
    pcm = plt.pcolormesh(xx, yy, n_difference, cmap=plt.get_cmap('Spectral_r'), shading='auto')
    # 添加标题
    ax.set_title('Ra={}, t={:.3f}'.format(Ra_star, t_diag), fontsize=16)
    # 添加颜色条
    #cbar = fig.colorbar(cs, orientation='vertical')
    # 设置坐标轴的比例相同，使得图形保持矩形
    ax.set_aspect('equal')
    # 添加颜色条，不需要指定boundaries
    cbar = fig.colorbar(pcm, ax=ax, orientation='vertical' ,fraction=0.02)
    cbar.set_label('T')
    plt.show()

if flag_pcolormesh_stream:
    t_diag=out_iter*dt
    fig, ax = plt.subplots(figsize=(10, 3))
    n_difference=n_density
    # 绘制伪彩色图
    
    plt.streamplot(xx, yy, u, v, linewidth=1)
    
    # plt.pcolormesh(xx, yy, n_difference, cmap=plt.get_cmap('Spectral_r'), vmin=-np.abs(n_difference).max(), vmax=np.abs(n_difference).max(), shading='auto')
    # plt.colorbar(label='n',boundaries=np.linspace(-np.abs(n_difference).max(), np.abs(n_difference).max()))
    
    # 使用自动缩放的颜色条
    pcm = plt.pcolormesh(xx, yy, n_difference-n_init0_matrix, cmap=plt.get_cmap('Spectral_r'), shading='auto')
    # 添加标题
    ax.set_title('Ra={}, t={}'.format(Ra_star, t_diag), fontsize=16)
    # 添加颜色条
    #cbar = fig.colorbar(cs, orientation='vertical')
    # 设置坐标轴的比例相同，使得图形保持矩形
    ax.set_aspect('equal')
    # 添加颜色条，不需要指定boundaries
    cbar = fig.colorbar(pcm, ax=ax, orientation='vertical' ,fraction=0.02)
    cbar.set_label('dT')
    plt.show()

v_slice = v[Ny//2, :]
n_slice=n[Ny//2, :]
heat_mid = np.dot(v_slice, n_slice)*dx
print(f'heat cross midplane:{heat_mid}')
# 创建图表
plt.figure(figsize=(10, 3))
plt.plot(x, v_slice)
plt.xlabel('x')
plt.ylabel('v-midplane')  # 这里假设y轴是某种值，您可以根据实际情况修改
plt.title('Ra={},t={:.3f}'.format(Ra,t_total))
plt.show()

if False:
    t_diag=out_iter*dt
    fig, ax = plt.subplots(figsize=(10, 3))
    n_difference=n_density
    # 绘制伪彩色图
    
    plt.streamplot(xx, yy, u, v, linewidth=1)
    
    # plt.pcolormesh(xx, yy, n_difference, cmap=plt.get_cmap('Spectral_r'), vmin=-np.abs(n_difference).max(), vmax=np.abs(n_difference).max(), shading='auto')
    # plt.colorbar(label='n',boundaries=np.linspace(-np.abs(n_difference).max(), np.abs(n_difference).max()))
    
    # 使用自动缩放的颜色条
    pcm = plt.pcolormesh(xx, yy, v, cmap=plt.get_cmap('Spectral_r'), shading='auto')
    # 添加标题
    ax.set_title('Ra={}, t={}'.format(Ra_star, t_diag), fontsize=16)
    # 添加颜色条
    #cbar = fig.colorbar(cs, orientation='vertical')
    # 设置坐标轴的比例相同，使得图形保持矩形
    ax.set_aspect('equal')
    # 添加颜色条，不需要指定boundaries
    cbar = fig.colorbar(pcm, ax=ax, orientation='vertical' ,fraction=0.02)
    cbar.set_label('v')
    plt.show()

#绘制流函数二维颜色图。
flag_contourf=False
if flag_contourf:
    fig, ax = plt.subplots(figsize=(10, 3))
    # 绘制等值线图，增加levels参数来使等值线更密集  
    levels = np.linspace(psi.min(), psi.max(), 20)  # 这里设置为50个级别  
    cs = ax.contourf(xx, yy, psi, levels=levels, cmap=plt.get_cmap('Spectral').reversed())  
    plt.streamplot(xx, yy, u, v, linewidth=0.5,color='black')
    # 添加标题  
    ax.set_title('Ra={},t={}'.format(Ra_star,t_total), fontsize=16)  # 设置标题文本和字体大小  
    ax.set_xlabel('X')  # 自行替换为合适的标签
    ax.set_ylabel('Y')  # 自行替换为合适的标签
    #添加colorbar
    cbar = fig.colorbar(cs)
    cbar.set_label('psi') 
    ax.set_aspect('equal')
    plt.show()

flag_contourf=False
if flag_contourf:
    fig, ax = plt.subplots()
    # 绘制等值线图，增加levels参数来使等值线更密集  
    levels = np.linspace(n.min(), n.max(), 30)  # 这里设置为50个级别  
    cs = ax.contourf(xx, yy, n, levels=levels, cmap=plt.get_cmap('Spectral').reversed())  
    # 添加标题  
    ax.set_title('Ra={},t={}'.format(Ra_star,t_total), fontsize=16)  # 设置标题文本和字体大小  
    # 添加x轴和y轴的标题
    ax.set_xlabel('X')  # 自行替换为合适的标签
    ax.set_ylabel('Y')  # 自行替换为合适的标签
    #添加colorbar
    cbar = fig.colorbar(cs)
    # 设置其标题
    cbar.set_label('T')  # 自行替换为合适的标签，例如'浓度'、'温度'等
    plt.show()


# fig, ax = plt.subplots()
# q = plt.quiver(xx, yy, u, v, scale=0.002)  # scale factor of 2 to make arrows larger
# ax.quiverkey(q, X=0.3, Y=1.1, U=1000,
#              label='Quiver key, length = 10', labelpos='E')
# plt.show()
